import numpy as np
import numexpr as ne
import pandas as pd
import re, os
import colorama
import warnings
warnings.filterwarnings("ignore")
# root_path=os.path.join(os.getcwd(),"Filter")
# TEXT_COLOR_RED = '\033[31m'
# TEXT_COLOR_GREEN = '\033[32m'
# TEXT_COLOR_YELLOW = '\033[33m'
# TEXT_COLOR_END = '\033[0m'

colorama.init(autoreset=True)
TEXT_COLOR_RED = colorama.Fore.RED
TEXT_COLOR_GREEN = colorama.Fore.GREEN
TEXT_COLOR_YELLOW = colorama.Fore.YELLOW
TEXT_COLOR_END = colorama.Fore.RESET
root_path=os.path.join(os.getcwd(),".\\Data_process\\CE_WAD")

def getCIE(RGB_pd, degree=None):
    if degree is None:
        x_key, y_key = 'CIEx', 'CIEy'
    else:
        x_key, y_key = 'CIEx_%d' % (degree,), 'CIEy_%d' % (degree,)
    xy = RGB_pd.loc[:, (x_key, y_key)].to_numpy()
    z = 1.0 - xy.sum(axis=1)
    return np.concatenate([xy, z[:, np.newaxis]], axis=1)


def getCE(RGB_pd, degree=None):
    if degree is None:
        ce_key = 'CE'
    else:
        ce_key = 'CE_%d' % (degree,)
    return RGB_pd.loc[:, ce_key].to_numpy()[:, np.newaxis]


def getAngles(RGB_pd):
    angles = []
    for col in RGB_pd.columns.values.tolist():
        matched = re.match('CIEx_([0-9]+)', col)
        if matched:
            angles.append(int(matched.groups()[0]))
    return angles


def w2rgb(W_CIE, R_CIE, G_CIE, B_CIE):
    if W_CIE.ndim == 2 and W_CIE.shape[0] == 1 and R_CIE.shape[0] > 1:
        W_CIE = W_CIE.repeat(R_CIE.shape[0], axis=0)
    r_xy = R_CIE[:, 0] / R_CIE[:, 1]
    r_zy = R_CIE[:, 2] / R_CIE[:, 1]
    g_xy = G_CIE[:, 0] / G_CIE[:, 1]
    g_zy = G_CIE[:, 2] / G_CIE[:, 1]
    b_xy = B_CIE[:, 0] / B_CIE[:, 1]
    b_zy = B_CIE[:, 2] / B_CIE[:, 1]
    w_xy = W_CIE[:, 0] / W_CIE[:, 1]
    w_zy = W_CIE[:, 2] / W_CIE[:, 1]
    A = r_xy - b_xy
    B = g_xy - b_xy
    C = w_xy - b_xy
    D = r_zy - b_zy
    E = g_zy - b_zy
    F = w_zy - b_zy
    R_L = 500 * (C * E - B * F) / (A * E - B * D)
    G_L = (500 * C - R_L * A) / B
    B_L = 500 - R_L - G_L
    RGB_L = np.concatenate([R_L[:, np.newaxis], G_L[:, np.newaxis], B_L[:, np.newaxis]], axis=1)
    return RGB_L


def rgb2w(RGB_L, R_CIE, G_CIE, B_CIE):
    if RGB_L.ndim == 2 and RGB_L.shape[0] == 1 and R_CIE.shape[0] > 1:
        RGB_L = RGB_L.repeat(R_CIE.shape[0], axis=0)
    r_xy = R_CIE[:, 0] / R_CIE[:, 1]
    r_zy = R_CIE[:, 2] / R_CIE[:, 1]
    g_xy = G_CIE[:, 0] / G_CIE[:, 1]
    g_zy = G_CIE[:, 2] / G_CIE[:, 1]
    b_xy = B_CIE[:, 0] / B_CIE[:, 1]
    b_zy = B_CIE[:, 2] / B_CIE[:, 1]
    X = RGB_L[:, 0] * r_xy + RGB_L[:, 1] * g_xy + RGB_L[:, 2] * b_xy
    Y = RGB_L[:, 0] + RGB_L[:, 1] + RGB_L[:, 2]
    Z = RGB_L[:, 0] * r_zy + RGB_L[:, 1] * g_zy + RGB_L[:, 2] * b_zy
    XYZ = np.concatenate([X[:, np.newaxis], Y[:, np.newaxis], Z[:, np.newaxis]], axis=1)
    return XYZ / XYZ.sum(axis=1)[:, np.newaxis]


def cie2uv(CIE):
    if CIE.ndim == 1:
        tmp = -2 * CIE[0] + 12 * CIE[1] + 3
        UV = np.array([4 * CIE[0], 9 * CIE[1]]) / tmp
    else:
        UV = np.empty((CIE.shape[0], 2))
        UV[:, 1] = -2 * CIE[:, 0] + 12 * CIE[:, 1] + 3
        UV[:, 0] = 4 * CIE[:, 0] / UV[:, 1]
        UV[:, 1] = 9 * CIE[:, 1] / UV[:, 1]
    return UV


def wce(RGB_L, R_CE, G_CE, B_CE):
    if RGB_L.ndim == 2 and RGB_L.shape[0] == 1 and R_CE.shape[0] > 1:
        RGB_L = RGB_L.repeat(R_CE.shape[0], axis=0)
    W_CE = RGB_L.sum(axis=1) / (RGB_L[:, 0] / R_CE[:, 0] + RGB_L[:, 1] / G_CE[:, 0] + RGB_L[:, 2] / B_CE[:, 0])
    return W_CE[:, np.newaxis]

def fout(input_file_name,process_name,out_format):
    info_list=input_file_name.split(".")[0].split("-")[0].split("_")
    fout_prefix = f'{process_name}_'+info_list[-1]
    ntry = 0
    while True:
        if ntry == 0:
            fout = fout_prefix + f'.{out_format}'
        else:
            fout = fout_prefix + '-%d.%s' % (ntry,out_format)
        if os.access(os.path.join(root_path,fout), os.F_OK):
            ntry += 1
        else:
            break
    return fout

def main(R_inp, G_inp, B_inp,config={},zf=False):
    print('Columns in R: %s' % (' '.join(R_inp.columns.values.tolist())))
    print('Columns in G: %s' % (' '.join(G_inp.columns.values.tolist())))
    print('Columns in B: %s' % (' '.join(B_inp.columns.values.tolist())))
    #匹配RGB
    def Match_RGB(R,G,B,match_keys):
        R_cond = R.loc[:, match_keys].to_numpy()
        G_cond = G.loc[:, match_keys].to_numpy()
        B_cond = B.loc[:, match_keys].to_numpy()

        diff_RG = np.empty((G_cond.shape[0],))
        diff_RB = np.empty((B_cond.shape[0],))

        RGB = []
        for Ri in range(R_cond.shape[0]):
            ref = R_cond[Ri, :]
            ne.evaluate('sum(abs(G_cond - ref), axis=1)', out=diff_RG)
            Gil = np.where(diff_RG < 1e-7)[0]
            if len(Gil) == 0:
                continue
            ne.evaluate('sum(abs(B_cond - ref), axis=1)', out=diff_RB)
            Bil = np.where(diff_RB < 1e-7)[0]
            if len(Bil) == 0:
                continue
            for Gi in Gil:
                for Bi in Bil:
                    RGB.append([Ri, Gi, Bi])
        RGB = np.array(RGB)
        try:
            R, G, B = R.loc[RGB[:, 0]], G.loc[RGB[:, 1]], B.loc[RGB[:, 2]]
        except IndexError:
            pass
        else:
            pass
        R.reset_index(drop=True)
        G.reset_index(drop=True)
        B.reset_index(drop=True)
        return RGB, R, G, B
    match_bool = True
    while match_bool:
        if zf:
            match_inp=config["match_inp"]
        else:
            match_inp = input('If you want to match the conditions of RGB, enter the variables to be matched: \n').strip()
            config["match_inp"]=match_inp
        match_keys = match_inp.strip().split()
        print(F"match_keys={match_keys}")
        if len(match_keys) > 0:
            RGB, R, G, B=Match_RGB(R_inp,G_inp,B_inp,match_keys)
            if len(RGB) == 0:
                print('0 Raws are Matched,Unable to match RGB with condition : %s '%(match_keys,))
                print(TEXT_COLOR_YELLOW + "ReMatching......" + TEXT_COLOR_END)
                continue
            MatchedRGB=R.loc[:, match_keys]
            # fout_mat=fout(R_inp,"Matched","csv")
            # MatchedRGB.to_csv(os.path.join(root_path,f"MatchedRGB_{suffix}.csv"), index=False)
            # print(f'Write output file',TEXT_COLOR_YELLOW+f'MatchedRGB_{suffix}.csv '+TEXT_COLOR_END,'successfully!' )
            if R.shape != G.shape or R.shape != B.shape:
                print('The shapes of RGB are not the same',"R: %d G: %d  B: %d"%(len(R),len(G),len(B)))
                print(TEXT_COLOR_YELLOW + "Need ReMatching......" + TEXT_COLOR_END)
                continue
            match_bool=False
        else:
            print("Input Can Not Be Empty! ")
            print(TEXT_COLOR_YELLOW + "ReMatching......" + TEXT_COLOR_END)
    ce_wad, col_names = [], []
    if match_keys is not None:
        match_cond = R.loc[:, match_keys].to_numpy()
        ce_wad.append(match_cond)
        col_names = match_keys.copy()

    #增添输出列
    if zf:
        addcond_inp=config['addcond_inp']
    else:
        addcond_inp = input( 'If you want to write additional conditions into the final output, enter the variables: \n').strip()
        config["addcond_inp"] = ""
    add_keys = addcond_inp.strip().split()
    if len(add_keys) > 0:
        ce_wad.append(R.loc[:, add_keys].to_numpy())
        ce_wad.append(G.loc[:, add_keys].to_numpy())
        ce_wad.append(B.loc[:, add_keys].to_numpy())
        for prefix in ('R_', 'G_', 'B_'):
            for key in add_keys:
                col_names.append(prefix + key)

    #输入白点及相关计算
    if zf:
        w_cie_xy=config["w_cie_xy"]
    else:
        w_cie_xy = input('Please input CIEx and CIEy of white point: \n').strip()
        config["w_cie_xy"] =w_cie_xy
    xy = [float(x) for x in w_cie_xy.split()]
    w_cie = np.array([xy[0], xy[1], 1.0 - xy[0] - xy[1]])[np.newaxis, :]
    angles = getAngles(R)
    nangle = len(angles)
    suffix_0 = None if nangle == 0 else 0  #以0度为参考
    r_cie, g_cie, b_cie = getCIE(R, suffix_0), getCIE(G, suffix_0), getCIE(B, suffix_0)
    r_ce, g_ce, b_ce = getCE(R, suffix_0), getCE(G, suffix_0), getCE(B, suffix_0)

    #读取或生成白光RGB参考文件
    if zf:
        ref=config["ref"]
    else:
        ref = input('Please input line number of ref (0 means no ref) or a csv that contain RGB_Lum and RGB_CE write by this script :\n').strip()
        config["ref"] = ref
    def rgb2w_df(ref):
        RGB2w_df_ref = pd.DataFrame()
        if isinstance(ref,pd.DataFrame):
            rgb2w_df=ref
            rgb_l_ref = rgb2w_df.loc[:, ['R_Lum', 'G_Lum', 'B_Lum']].to_numpy()
            rgb_ce_ref = rgb2w_df.loc[:, ['R_CE', 'G_CE', 'B_CE']].to_numpy()
        elif ref.isdigit():
            ref = int(ref)
            if ref > 1:
                rgb_l_ref = w2rgb(w_cie, r_cie[ref - 2][np.newaxis, :], g_cie[ref - 2][np.newaxis, :],
                                  b_cie[ref - 2][np.newaxis, :])
                rgb_ce_ref = np.array([r_ce[ref - 2, 0], g_ce[ref - 2, 0], b_ce[ref - 2, 0]])[np.newaxis, :]
            else:
                rgb_l_ref = w2rgb(w_cie, r_cie, g_cie, b_cie)
                rgb_ce_ref = np.concatenate([r_ce, g_ce, b_ce], axis=1)

            RGB2w_df_ref = pd.DataFrame(data=np.concatenate([rgb_l_ref, rgb_ce_ref], axis=1),
                                    columns=['R_Lum', 'G_Lum', 'B_Lum', 'R_CE', 'G_CE', 'B_CE'])
        else:
            rgb2w_df = pd.read_csv(os.path.join(root_path,ref))
            rgb_l_ref = rgb2w_df.loc[:, ['R_Lum', 'G_Lum', 'B_Lum']].to_numpy()
            rgb_ce_ref = rgb2w_df.loc[:, ['R_CE', 'G_CE', 'B_CE']].to_numpy()

        return rgb_l_ref, rgb_ce_ref ,RGB2w_df_ref
    rgb_l_ref, rgb_ce_ref,RGB2w_df_ref=rgb2w_df(ref)
    if rgb_l_ref.ndim == 2 and rgb_l_ref.shape[0] == 1:
        rgb_l_ref = rgb_l_ref.repeat(r_cie.shape[0], axis=0)
        rgb_ce_ref = rgb_ce_ref.repeat(r_cie.shape[0], axis=0)

    ####计算白光####
    w_ce = wce(rgb_l_ref, r_ce, g_ce, b_ce)
    ce_wad.append(w_ce)
    col_names.append('CE')
    w_uv = cie2uv(w_cie)

    #1 RGB部分的L-Decay
    # 重写：因为无法改动已经写好的JNCD部分。
    rgb_lum_bool = True if (angles and angles[0]) == 0 else False
    rgb_lum = []
    rgb_lum_col_names = []
    if nangle > 0:
        for color_name, color, ce in zip(['R_', 'G_', 'B_'], [R, G, B], [r_ce, g_ce, b_ce]):
            for i, angle in enumerate(angles):
                if rgb_lum_bool:
                    rgb_lum.append(getCE(color, angle) / ce)
                    rgb_lum_col_names.append(color_name + str(angle) + '_Lum')

    #2 W部分的L-Decay，理论上应该复用JNCD部分rgb_ang_l的计算
    w_lum_bool = True if (angles and angles[0] == 0 )else False
    w_lum = []
    w_lum_col_names = []
    if nangle > 0:
        rgb_l_0 = w2rgb(w_cie, r_cie, g_cie, b_cie)
        rgb_ce_0 = np.concatenate([r_ce, g_ce, b_ce], axis=1)
        w_ang_l = (rgb_l_0[:, 0] * getCE(R, 0)[:, 0] / rgb_ce_0[:, 0]) + \
                  (rgb_l_0[:, 1] * getCE(G, 0)[:, 0] / rgb_ce_0[:, 1]) + \
                  (rgb_l_0[:, 2] * getCE(B, 0)[:, 0] / rgb_ce_0[:, 2])
        for i, angle in enumerate(angles):
            i_w_lum = ((rgb_l_0[:, 0] * getCE(R, angle)[:, 0] / rgb_ce_0[:, 0]) + \
                       (rgb_l_0[:, 1] * getCE(G, angle)[:, 0] / rgb_ce_0[:, 1]) + \
                       (rgb_l_0[:, 2] * getCE(B, angle)[:, 0] / rgb_ce_0[:, 2])) / w_ang_l
            w_lum.append(i_w_lum[:, np.newaxis])
            w_lum_col_names.append('W_%d' % (angle) + '_Lum')

    #3 计算U、V
    uv_bool = True if (angles and angles[0] == 0 )else False
    jncd_bool = True if (angles and angles[0] == 0 ) else False
    if nangle > 0:
        rgb_ang_l = np.empty(rgb_l_ref.shape, dtype=np.double)
        relative_uv = np.empty((rgb_l_ref.shape[0], 2), dtype=np.double)
        old_uv_columns = []
        for i, angle in enumerate(angles):
            rgb_ang_l[:, 0] = rgb_l_ref[:, 0] * getCE(R, angle)[:, 0] / rgb_ce_ref[:, 0]
            rgb_ang_l[:, 1] = rgb_l_ref[:, 1] * getCE(G, angle)[:, 0] / rgb_ce_ref[:, 1]
            rgb_ang_l[:, 2] = rgb_l_ref[:, 2] * getCE(B, angle)[:, 0] / rgb_ce_ref[:, 2]
            r_ang_cie, g_ang_cie, b_ang_cie = getCIE(R, angle), getCIE(G, angle), getCIE(B, angle)
            w_ang_cie = rgb2w(rgb_ang_l, r_ang_cie, g_ang_cie, b_ang_cie)
            w_ang_uv = cie2uv(w_ang_cie)
            ce_wad.append(w_ang_uv)
            col_names.extend(['U_%d' % (angle), 'V_%d' % (angle)])
            old_uv_columns.extend(['U_%d' % (angle), 'V_%d' % (angle)])
            if jncd_bool:
                if i == 0:
                    uv_ref = w_ang_uv
                    relative_uv = w_ang_uv
                    abs_jncd = []
                    abs_jncd_col_names = []
                    rel_jncd = []
                    rel_jncd_col_names = []
                else:
                    abs_jncd.append(np.sqrt(np.square(w_ang_uv - uv_ref).sum(axis=1))[:, np.newaxis] * 250)
                    abs_jncd_col_names.append('ABS_JNCD_%d' % (angle,))
                    rel_jncd.append(np.sqrt(np.square(w_ang_uv - relative_uv).sum(axis=1))[:, np.newaxis] * 250)
                    rel_jncd_col_names.append('REL_JNCD_%d' % (angle,))
                    relative_uv = w_ang_uv
        if jncd_bool:
            ce_wad.extend(abs_jncd)
            ce_wad.extend(rel_jncd)
            col_names.extend(abs_jncd_col_names)
            col_names.extend(rel_jncd_col_names)
    else:
        rgb_cur_l = np.empty(rgb_l_ref.shape, dtype=np.double)
        rgb_cur_l[:, 0] = rgb_l_ref[:, 0] * r_ce[:, 0] / rgb_ce_ref[:, 0]
        rgb_cur_l[:, 1] = rgb_l_ref[:, 1] * g_ce[:, 0] / rgb_ce_ref[:, 1]
        rgb_cur_l[:, 2] = rgb_l_ref[:, 2] * b_ce[:, 0] / rgb_ce_ref[:, 2]
        w_ang_cie = rgb2w(rgb_cur_l, r_cie, g_cie, b_cie)
        w_ang_uv = cie2uv(w_ang_cie)
        ce_wad.append(w_ang_uv)
        col_names.extend(['U', 'V'])

    #4 将以上计算结果拼接到最终表
    if w_lum_bool:
        ce_wad.extend(w_lum)
        col_names.extend(w_lum_col_names)
    if rgb_lum_bool:
        ce_wad.extend(rgb_lum)
        col_names.extend(rgb_lum_col_names)

    #5 拼接结果
    ce_wad = np.concatenate(ce_wad, axis=1)
    ce_wad_df = pd.DataFrame(data=ce_wad, columns=col_names)

    #6 选择要返回的结果信息
    if uv_bool:
        uv_columns = ['U_' + str(angle) for angle in angles] + ['V_' + str(angle) for angle in angles]
    else:
        uv_columns=['U','V']
    old_col_names_list = ce_wad_df.columns.to_list()
    start_index = old_col_names_list.index(uv_columns[0])
    end_index = old_col_names_list.index(uv_columns[-1]) + 1
    new_col_names_list = old_col_names_list[:start_index] + uv_columns + old_col_names_list[end_index:]
    ce_wad_df = ce_wad_df[new_col_names_list]

    #7 处理结果列
    if w_lum_bool:
        ce_wad_df[w_lum_col_names] = ce_wad_df[w_lum_col_names].applymap(lambda x: str(round(x * 100, 2)) + '%')
    if rgb_lum_bool:
        ce_wad_df[rgb_lum_col_names] = ce_wad_df[rgb_lum_col_names].applymap(lambda x: str(round(x * 100, 2)) + '%')

    return RGB2w_df_ref,MatchedRGB,ce_wad_df,config

if __name__ == '__main__':
    root_path=os.getcwd()
    rgb_inp = input('Please input the file of R, G, and B: \n').strip()
    R, G, B = rgb_inp.split()
    R_finp, G_finp, B_finp = pd.read_csv(os.path.join(root_path,R)), pd.read_csv(os.path.join(root_path,G)), pd.read_csv(os.path.join(root_path,B))
    rgb2w_df_ref,Matchedrgb,ce_wad_df,config=main(R_finp, G_finp, B_finp,config={},zf=False)
    fout_rgb,fout_mat,fout_ce =fout(R, "RGB2W", "csv"), fout(R, "MatchedRGB", "csv"),fout(R, "CE_WAD", "csv")
    rgb2w_df_ref.to_csv(os.path.join(root_path, fout_rgb), index=False)
    Matchedrgb.to_csv(os.path.join(root_path,fout_mat), index=False)
    ce_wad_df.to_csv(os.path.join(root_path, fout_ce), index=False)
    print( 'Write output file'+TEXT_COLOR_GREEN+" %s,%s,%s "% (fout_mat,fout_rgb,fout_ce)+TEXT_COLOR_END+' successfully!' )


'''
Test_Iput:
R_5.csv G_5.csv B_5.csv # R-F.csv G-F.csv Filtered_B_Ref.csv # R.csv G.csv B.csv
SiON-i.d	SiON-B.d    # SiON-B.d	SiON-i2.d     # SiON-i.d	SiON-B.d
CE	CIEx	CIEy        #CE_0	CIEx_0	CIEy_0     # CE_0	CIEx_0	CIEy_0
0.3 0.315
0
'''
# Filter_R_Ref.csv  Filter_G_Ref.csv  Filter_B_Ref.csv
# Matched_R_5.csv Matched_G_5.csv Matched_B_5.csv




# Matched_B_5.csv Matched_G_5.csv Matched_R_5.csv
#CE,CIEx,CIEy
# Matched_B_5.csv Matched_G_5.csv  Matched_R_5.csv
# Filtered_R_Ref.csv Filtered_G_Ref.csv Filtered_B_Ref.csv
# Filter_R_Ref-2.csv Filter_G_Ref-2.csv Filter_B_Ref-2.csv